package com.corn.flink.lesson5;

import com.corn.flink.lesson4.Event;
import org.apache.flink.api.common.eventtime.SerializableTimestampAssigner;
import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.api.common.functions.AggregateFunction;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.windowing.ProcessWindowFunction;
import org.apache.flink.streaming.api.windowing.assigners.SlidingEventTimeWindows;
import org.apache.flink.streaming.api.windowing.assigners.TumblingEventTimeWindows;
import org.apache.flink.streaming.api.windowing.time.Time;
import org.apache.flink.streaming.api.windowing.windows.TimeWindow;
import org.apache.flink.util.Collector;

import java.sql.Timestamp;
import java.time.Duration;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;
import java.util.concurrent.TimeUnit;

/**
 * @author JimWu
 * @date 2023/3/15 16:52
 * 基于第一个参数（增量聚合函数）来处理窗口数据，每来一个数
 * 据就做一次聚合；等到窗口需要触发计算时，则调用第二个参数（全窗口函数）的处理逻辑输
 * 出结果。需要注意的是，这里的全窗口函数就不再缓存所有数据了，而是直接将增量聚合函数
 * 的结果拿来当作了 Iterable 类型的输入。一般情况下，这时的可迭代集合中就只有一个元素了。
 **/
public class WindowAggregateFunctionWithProcessFunction {

    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(1);
        DataStreamSource<Event> ds = env.fromElements(
                new Event("add", 1000),
                new Event("delete", 2000),
                new Event("update", 3000),
                new Event("select", 5000),
                new Event("add", 6000),
                new Event("delete", 7000),
                new Event("select", 8000),
                new Event("select", 9000),
                new Event("select", 10000),
                new Event("select", 11000)
        );


        ds.assignTimestampsAndWatermarks(
                WatermarkStrategy.<Event>forBoundedOutOfOrderness(Duration.ZERO)
                        .withTimestampAssigner(new SerializableTimestampAssigner<Event>() {
                            @Override
                            public long extractTimestamp(Event element, long recordTimestamp) {
                                return element.timestamp;
                            }
                        }))
                .keyBy(data -> data.type)
                .window(SlidingEventTimeWindows.of(Time.seconds(5), Time.seconds(2)))
                .aggregate(new AggregateFunction<Event, Map<String, Integer>, Map<String, Integer>>() {
                    @Override
                    public Map<String, Integer> createAccumulator() {
                        return new HashMap<>();
                    }

                    @Override
                    public Map<String, Integer> add(Event value, Map<String, Integer> accumulator) {
                        if (accumulator.get(value.getType()) == null) {
                            accumulator.put(value.type, 1);
                        } else {
                            accumulator.put(value.type, accumulator.get(value.type) + 1);
                        }

                        return accumulator;
                    }

                    @Override
                    public Map<String, Integer> getResult(Map<String, Integer> accumulator) {
                        return accumulator;
                    }

                    @Override
                    public Map<String, Integer> merge(Map<String, Integer> a, Map<String, Integer> b) {
                        return null;
                    }
                }, new ProcessWindowFunction<Map<String, Integer>, String, String, TimeWindow>() {
                    @Override
                    public void process(String s, Context context, Iterable<Map<String, Integer>> elements, Collector<String> out) throws Exception {
                        Timestamp startTime = new Timestamp(context.window().getStart());
                        Timestamp endTime = new Timestamp(context.window().getEnd());
                        Iterator<Map<String, Integer>> it = elements.iterator();
                        Map<String, Integer> next = it.next();
                        next.keySet().forEach(key -> {
                            out.collect(startTime + " - " + endTime + " -> " + key + " " + next.get(key));
                        });
                    }
                })
                .print();

        env.execute();
    }
}
